package com.study.basic

import org.apache.spark.mllib.linalg
import org.apache.spark.mllib.stat.Statistics
import org.apache.spark.sql.SparkSession

/**
 * 相关性
 *
 * @author stephen
 * @date 2019-08-27 11:10
 */
object CorrelationDemo {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName(this.getClass.getSimpleName)
      .master("local[*]")
      .getOrCreate()

    spark.sparkContext.setLogLevel("warn")

    val data = spark.sparkContext.parallelize(
      Seq(
        linalg.Vectors.dense(-1, 2, 3),
        linalg.Vectors.dense(2, -4, 6),
        linalg.Vectors.dense(3, 6, -9)
      ))

    // 初始化数据
    val seriesX = spark.sparkContext.parallelize(Seq(1.0, 2.0, 3.0, 4.0))
    val seriesY = spark.sparkContext.parallelize(Seq(2.0, 3.0, 5.0, 7.0))
    // 皮尔逊计算相关性
    val pearsonCorrelation = Statistics.corr(seriesX, seriesY, "pearson")
    println(s"pearson correlation: ${pearsonCorrelation}")

    // 斯皮尔曼计算相关性
    val spearmanCorrelation = Statistics.corr(seriesX, seriesY, "spearman")
    println(s"spearman correlation: ${spearmanCorrelation}")

    // 皮尔逊计算相关系数矩阵,由矩阵各列间的相关系数构成
    val matrixCorrelation = Statistics.corr(data, "pearson")
    println(s"matrix pearson correlation:\n${matrixCorrelation}")

    spark.stop()
  }
}
